import tushare as ts
from datetime import time
from datetime import date, datetime, timedelta
from vnpy.trader.constant import Status, Interval, Direction
import pandas as pd

from vnpy_ctastrategy import (
    CtaTemplate,
    StopOrder,
    TickData,
    BarData,
    TradeData,
    OrderData,
    BarGenerator,
    ArrayManager,
)


class IndexDrivenStrategy(CtaTemplate):
    """基于中证1000指数信号的个股交易策略"""

    author = "Youpeng Hu"
    
    # 策略参数

    index_symbol = '000852.SH'
    rsi_window = 14
    signal_lead = 1
    trail_percent = 0.8

    parameters = [
        "index_symbol",     # 中证1000指数代码
        "rsi_window",                 # RSI计算窗口
        "signal_lead",                # 信号提前天数
        "trail_percent"              # 动态跟踪止盈比例
    ]

    def __init__(self, cta_engine, strategy_name, vt_symbol, setting):
        super().__init__(cta_engine, strategy_name, vt_symbol, setting)
        self.data_cache = {}  # 日期: 数据字典
        self.current_date = None
        self.index_rsi = 0
        self.signal_queue = {}
        
        # 设置本地数据参数
        self.csv_path = "E:\\GitCloneProject\\quant25\\01Src\\vnpy\\environment_evaluation\\data\\zz1000_data.csv"
        self.rsi_window = setting.get("rsi_window", 14)

        # 初始化BarGenerator和ArrayManager
        self.bg = BarGenerator(self.on_bar)
        self.am = ArrayManager()

    def on_init(self):
        """初始化策略"""
        self.write_log("策略初始化")
        self.load_historical_data()  # 加载本地历史数据
        self.rsi_buy = 50 + (100 - 30)  # 动态计算超买/超卖阈值
        self.rsi_sell = 50 - (100 - 70)

    def load_historical_data(self):
        """从本地CSV加载指数历史数据"""

        try:
            df = pd.read_csv(self.csv_path, parse_dates=['datetime'])
        except Exception as e:
            self.write_log(f"读取数据失败：{str(e)}")
            raise

        # 数据标准化处理
        required_columns = ['datetime', 'close', 'open', 'high', 'low']
        if not all(col in df.columns for col in required_columns):
            missing = [col for col in required_columns if col not in df.columns]
            self.write_log(f"缺失必要字段：{missing}")
            raise ValueError(f"数据文件缺少必要字段：{missing}")

        # 转换数据类型
        df['close'] = df['close'].astype(float)
        df['open'] = df['open'].astype(float)
        df['high'] = df['high'].astype(float)
        df['low'] = df['low'].astype(float)

        # 按日期排序并设置索引
        df.rename(columns={'datetime': 'date'}, inplace=True)
        print(df.columns)

        # 存入缓存（使用日期字符串作为键）
        self.data_cache = df.to_dict('dict')

        self.write_log(f"成功加载{len(df)}条数据")

    def get_index_data(self, date_str):
        """获取指定日期的指数数据"""
        date = pd.to_datetime(date_str)
        if date not in self.data_cache:
            self.write_log(f"未找到{date}的数据")
            return None
        return self.data_cache[date]

    def calculate_rsi(self, prices):
        """改进型RSI计算（包含异常值处理）"""
        delta = prices.diff(1)
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)
        
        avg_gain = gain.rolling(self.rsi_window, min_periods=1).mean()
        avg_loss = loss.rolling(self.rsi_window, min_periods=1).mean()
        
        rs = avg_gain / (avg_loss + 1e-8)  # 防止除零错误
        rsi = 100 - (100 / (1 + rs))
        return rsi

    def generate_signals(self):
        """生成交易信号"""
        if not self.data_cache:
            self.write_log("数据缓存为空，无法生成信号")
            return

        # 获取最新日期数据
        latest_date = max(self.data_cache.keys())
        index_data = self.get_index_data(latest_date)
        
        if index_data is None:
            return
        
        # 计算RSI指标
        close_prices = [self.get_index_data(d)['close'] for d in 
                        sorted(self.data_cache.keys()) if d <= latest_date]
        if len(close_prices) < self.rsi_window:
            self.write_log(f"数据不足，需至少{self.rsi_window}天数据")
            return
        
        self.index_rsi = self.calculate_rsi(pd.Series(close_prices))
        
        # 生成信号（考虑提前量）
        signal_date = latest_date - pd.Timedelta(days=self.signal_lead)
        if self.index_rsi < 30:
            self.signal_queue.setdefault(signal_date, []).append(('buy', self.symbol))
        elif self.index_rsi > 70:
            self.signal_queue.setdefault(signal_date, []).append(('sell', self.symbol))

    def process_signals(self):
        """处理待执行信号"""
        today = pd.to_datetime(self.bg.bar.datetime[-1] if self.bg.bar else datetime.now())
        today_str = today.strftime("%Y%m%d")
        
        if today_str in self.signal_queue:
            for action, symbol in self.signal_queue[today_str]:
                if action == 'buy' and self.pos == 0:
                    self.buy(price=self.get_current_price(), size=self.fixed_size)
                elif action == 'sell' and self.pos > 0:
                    self.sell(price=self.get_current_price(), size=self.fixed_size, stop=True)
            del self.signal_queue[today_str]

    def get_current_price(self):
        """获取个股当前价格"""
        return self.am.close[-1] if self.am.inited else 0.0

    def on_bar(self, bar: BarData):
        """Bar数据更新回调"""
        self.bg.update_tick(bar)
        self.cancel_all()

        am = self.am
        am.update_bar(bar)
        if not am.inited:
            return

        # 更新个股持仓相关指标
        self.intra_trade_high = bar.high_price
        self.intra_trade_low = bar.low_price

        # 处理指数信号
        self.generate_signals()
        self.process_signals()

        # 执行个股交易逻辑
        if self.pos == 0:
            if self.index_rsi < 30:
                self.buy(bar.close_price + 5, self.fixed_size)
        elif self.pos > 0:
            trailing_stop = self.intra_trade_high * (1 - self.trail_percent / 100)
            self.sell(trailing_stop, abs(self.pos), stop=True)

        self.put_event()

    def on_order(self, order: OrderData):
        """订单更新回调"""
        pass

    def on_trade(self, trade: TradeData):
        """交易更新回调"""
        self.put_event()

    def on_stop_order(self, stop_order: StopOrder):
        """止损单更新回调"""
        pass